'''
Author: goog
Date: 2022-01-07 20:24:27
LastEditTime: 2022-01-07 20:29:58
LastEditors: goog
Description: 
FilePath: /chengdu/TensorRT20220104/Models/classification/mobilenet/cls.py
Time Limit Exceeded!
'''
from .model import mobilenet_v3_large
import torch

def _cls(n_class, weights, device):
    model = mobilenet_v3_large(num_classes=n_class)
    device = torch.device(device)
    state_dict = torch.load(weights, map_location=device)
    model.load_state_dict(state_dict, strict=False)
    model.to(device)
    model.eval()
    return model
